import multiprocessing
import os
import subprocess
import sys
from typing import Dict, Generator, Union, List, Any


def run_script_linux(
        workspace: str,
        exec_user: str,
        interpreter: str,
        script: List[str]
) -> Generator[Union[Dict[str, Any], bytes], None, None]:
    """Execute a script locally in a pseudo-terminal and stream the output.

    Args:
        workspace: Directory to run the script in.
        exec_user: User to execute the script as.
        interpreter: The interpreter to use (e.g., '/bin/bash').
        script: List of command arguments (including the script path).

    Yields:
        Dictionary with initial process info, then bytes of output,
        finally a dictionary with exit status.
    """
    import pty
    import pwd

    pid, fd = pty.fork()
    if pid == 0:  # Child process
        try:
            if workspace:
                os.chdir(workspace)
            if exec_user:
                os.setuid(pwd.getpwnam(exec_user).pw_uid)
            os.execvp(interpreter, [interpreter, script])
        except Exception as e:
            print(f"Child process failed: {e}", file=sys.stderr)
            os._exit(1)

    # Parent process
    yield {'end': False, 'pid': pid}

    history = b''
    try:
        while True:
            try:
                data = os.read(fd, 1024)
                if not data:  # EOF
                    break
                history += data
                yield data
            except (OSError, IOError):
                break
    finally:
        os.close(fd)
        _, exit_status = os.waitpid(pid, 0)
        yield {
            'end': True,
            'exit_status': exit_status,
            'history': history
        }


def run_script_windows(
        workspace: str,
        interpreter: str,
        script: List[str],
        exec_user: str = None  # Windows 可能需要额外处理来切换用户
) -> Generator[Union[Dict[str, Any], bytes], None, None]:
    """Execute a script locally on Windows and stream the output.

    Args:
        workspace: Directory to run the script in.
        interpreter: The interpreter to use (e.g., 'python.exe').
        script: List of command arguments (including the script path).
        exec_user: Not fully supported on Windows (would require runas).

    Yields:
        Dictionary with initial process info, then bytes of output,
        finally a dictionary with exit status.
    """
    if exec_user:
        sys.stderr.write("Warning: Switching users is not fully supported on Windows\n")

    # 准备命令
    command = [interpreter, script]

    # 设置进程创建标志以创建新控制台
    creationflags = subprocess.CREATE_NEW_CONSOLE

    # 启动进程
    process = subprocess.Popen(
        command,
        cwd=workspace if workspace else None,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        stdin=subprocess.PIPE,
        creationflags=creationflags,
        text=False,  # 获取字节输出
        bufsize=1  # 行缓冲
    )

    yield {'end': False, 'pid': process.pid}

    history = b''
    try:
        while True:
            # 读取输出
            output = process.stdout.read(1024)  # read1 更高效
            if not output and process.poll() is not None:
                break
            if output:
                history += output
                yield output
    finally:
        # 获取退出状态
        exit_status = process.wait()

        yield {
            'end': True,
            'exit_status': exit_status,
            'history': history
        }
